/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 */

package com.huawei.housekeeper.common;

import static com.huawei.housekeeper.common.constant.SecurityConstants.DEFAULT_KEY_SIZE;
import static com.huawei.housekeeper.common.constant.SecurityConstants.ISSUER;
import static com.huawei.housekeeper.common.constant.SecurityConstants.KEYID;
import static com.huawei.housekeeper.common.constant.SecurityConstants.OIDC_ALG;
import static com.huawei.housekeeper.common.constant.SecurityConstants.USE_SIG;
import static org.junit.jupiter.api.Assertions.assertNotNull;

import com.huawei.housekeeper.common.constant.CommonConstants;
import com.huawei.housekeeper.common.utils.JwkUtil;

import com.google.common.base.Charsets;
import com.google.common.io.CharStreams;
import com.nimbusds.jose.jwk.JWK;
import com.nimbusds.jose.jwk.JWKSet;

import io.jsonwebtoken.Claims;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.SignatureAlgorithm;
import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.time.DateUtils;
import org.jose4j.jwk.JsonWebKey;
import org.jose4j.jwk.PublicJsonWebKey;
import org.jose4j.jwk.RsaJsonWebKey;
import org.jose4j.jwk.RsaJwkGenerator;
import org.json.JSONArray;
import org.json.JSONObject;

import java.io.InputStream;
import java.io.InputStreamReader;
import java.security.PrivateKey;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * JWK生成测试类
 *
 * @author Y84241604
 * @since 2022-06-02
 */
@Slf4j
public class SaasHousekeeperJwkTest {
    public static void main(String[] args) throws Exception {
        TestCreateKeystore();
    }

    /**
     * 根据kid生成jwk
     *
     * @throws Exception
     */

    static void TestCreateKeystore() throws Exception {
        // 根据kid生成jwk,用文件keystore.jwks存储
        RsaJsonWebKey jwk = RsaJwkGenerator.generateJwk(DEFAULT_KEY_SIZE);
        // sig or enc
        jwk.setUse(USE_SIG);
        jwk.setKeyId(KEYID);
        jwk.setAlgorithm(OIDC_ALG);
        final String publicKeyString = jwk.toJson(JsonWebKey.OutputControlLevel.PUBLIC_ONLY);
        final String privateKeyString = jwk.toJson(JsonWebKey.OutputControlLevel.INCLUDE_PRIVATE);

        assertNotNull(publicKeyString);
        assertNotNull(privateKeyString);
//        log.info("PublicKey: " + publicKeyString);
//        log.info("PrivateKey: " + privateKeyString);

        // 生成 jwks
        JSONObject jsonObject = new JSONObject(jwk.toParams(JsonWebKey.OutputControlLevel.INCLUDE_PRIVATE));
//        log.info("jwk: " + jsonObject);
        JSONArray jsonArray = new JSONArray();
        jsonArray.put(jsonObject);
        JSONObject jwksJSON = new JSONObject();
        jwksJSON.put("keys", jsonArray);
        log.info("jwks: " + jwksJSON);

        // 生成JWT
        PrivateKey privateKey = RsaJsonWebKey.Factory.newPublicJwk(privateKeyString).getPrivateKey();
        Map<String, Object> claims = new HashMap<>();
        claims.put("name", "admin");
        claims.put("password", "lhKk101@mm.");
        final Date exp = DateUtils.addMinutes(new Date(), 60);
        String jwt = Jwts.builder()
            .setClaims(claims)
            .setExpiration(exp)
            .setIssuer(ISSUER)
            .setSubject(ISSUER)
            .signWith(SignatureAlgorithm.RS256, privateKey)
            .compact();
//        log.info("jwt: " + jwt);
    }

    /**
     * 读取jwk,获取PrivateKeyString
     *
     * @throws Exception
     */
    static void TestGetPrivateKey() throws Exception {
        InputStream is = SaasHousekeeperJwkTest.class.getClassLoader().getResourceAsStream("keystore.jwks");
        String jwkString = CharStreams.toString(new InputStreamReader(is, Charsets.UTF_8));
        JWKSet jwkSet = JWKSet.parse(jwkString);
        assertNotNull(jwkSet);
        List<JWK> keys = jwkSet.getKeys();
        for (JWK key : keys) {
            PublicJsonWebKey jwk = RsaJsonWebKey.Factory.newPublicJwk(key.toString());
            final String publicKeyString = jwk.toJson(JsonWebKey.OutputControlLevel.PUBLIC_ONLY);
            final String privateKeyString = jwk.toJson(JsonWebKey.OutputControlLevel.INCLUDE_PRIVATE);
            log.info("publicKey : " + publicKeyString);
            log.info("privateKey : " + privateKeyString);
        }
        Map<String, Object> claims = new HashMap<>();
        claims.put(CommonConstants.User.USER_NAME, "Admin");
        claims.put(CommonConstants.User.USER_PAW, "lhKk101@mm.");
        String jwt = Jwts.builder()
            .setClaims(claims)
            .setExpiration(new Date(System.currentTimeMillis() + 3600 * 1000))
            .setIssuer(ISSUER)
            .setSubject(ISSUER)
            .signWith(SignatureAlgorithm.RS256, JwkUtil.getPrivateKeyByJwk())
            .compact();
        log.info("jwt: " + jwt);

        String token
            = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c";
        Claims claim2 = Jwts.parser()
            .setSigningKey(JwkUtil.getPublicKeyByJwk())
            .parseClaimsJws(token.replace("Bearer", "").replace(" ", ""))
            .getBody();
        log.info("claim2: " + claim2);
    }

}
